'''
@Company: TWL
@Author: xue jian
@Email: xuejian@kanzhun.com
@Date: 2020-04-09 19:42:20
'''
from multiprocessing import Process, Queue
import os, time, random
import tensorflow as tf

def write(q):
    a = [random.randint(0, 5) for i in range(50)]
    t = 0
    while t < len(a):
        if not q.full():
            q.put(a[t])
            print('put: ', a[t])
            time.sleep(0.05)
            t += 1

# def write():
#     a = [random.randint(0, 5) for i in range(50)]
#     return a

def init_tf():
    data_ph = tf.placeholder(dtype=tf.int64, shape=(None,), name='label')
    cate_tmp_column = tf.feature_column.categorical_column_with_vocabulary_list('fea', [0, 1, 2, 3, 4],
                                                                                     dtype=tf.int64, default_value=-1)
    cate_tmp_emb_column = tf.feature_column.embedding_column(cate_tmp_column, 3)
    emb = tf.feature_column.input_layer({'fea': data_ph}, [cate_tmp_emb_column])  # shape=[B, emb_size]
    sess = tf.Session()
    sess.run(tf.global_variables_initializer())
    sess.run(tf.tables_initializer())
    return sess, emb, data_ph

def train(d, sess, emb, data_ph):
    print("train! data = ", d)
    sess.run(emb, feed_dict={data_ph: d})
    print("train! finish ", d)
# 读数据进程执行的代码:
def read(q):
    sess, emb, data_ph = init_tf()
    batch_data = []
    while True:
        tmp_d = q.get()
        print("get: ", tmp_d)
        batch_data.append(tmp_d)
        if len(batch_data) >= 5:
            train(batch_data, sess, emb, data_ph)
            batch_data = []


# def train2(d):
#     print("22train! data = ", d)
#     time.sleep(1.0)
#     print("22train! finish ", d)
# # 读数据进程执行的代码:
# def read2(q):
#     batch_data = []
#     while True:
#         tmp_d = q.get()
#         print("22get: ", tmp_d)
#         batch_data.append(tmp_d)
#         if len(batch_data) >= 5:
#             train2(batch_data)
#             batch_data = []

import numpy as np
if __name__=='__main__':
    # 父进程创建Queue，并传给各个子进程：
    q = Queue(15)
    pw = Process(target=write, args=(q,))
    pr = Process(target=read, args=(q, ))
    # pr2 = Process(target=read2, args=(q,))
    # 启动子进程pw，写入:
    pw.start()
    # 启动子进程pr，读取:
    pr.start()
    # pr2.start()
    # 等待pw结束:
    pw.join()
    # pr进程里是死循环，无法等待其结束，只能强行终止:
    time.sleep(5)
    pr.terminate()
    # pr2.terminate()
    #
    # # t = write()
    # # train(t[:5], sess, emb, data_ph)